
#include "GPU_CRBM_MNIST.h"

using namespace loirey::fft;

int main(int argc, const char* argv[])
{
	int nRetCode = 0;
	//CMyRandom::InitRandomSeed();
	CImageEx::Initialize();

	if (1)
	{
		string strFN_Configuration		=	argv[1];
		string strOperation				=	argv[2];

		strOperation += "00000000";

		int op_TrainCDBN				=	CDataStructureTools::Hex2Num(strOperation[0]);
		int op_DrawCDBN					=	CDataStructureTools::Hex2Num(strOperation[1]);

		char strTemp[10001] = "";

		//////////	Config Loading	//////////

		CConfiguration Config(strFN_Configuration);
		Config.Load();

		string strFN_OutLog =
			Config.GetString("strFN_OutLog", "Global.log");
		string strFN_ShortLog =
			Config.GetString("strFN_ShortLog", "Short.log");

		sprintf_s(strTemp, "del %s", strFN_OutLog.c_str());
		CExternalTools::ConsoleTools::RunDosCommand("", strTemp);
		sprintf_s(strTemp, "del %s", strFN_ShortLog.c_str());
		CExternalTools::ConsoleTools::RunDosCommand("", strTemp);

		EnvVar.strFN_MNIST_Image_Train =
			Config.GetString("strFN_MNIST_Image_Train");
		EnvVar.strFN_MNIST_Label_Train =
			Config.GetString("strFN_MNIST_Label_Train");
		EnvVar.strFN_MNIST_Image_Test =
			Config.GetString("strFN_MNIST_Image_Test");
		EnvVar.strFN_MNIST_Label_Test =
			Config.GetString("strFN_MNIST_Label_Test");

		EnvVar.strFN_CDBN_Settings = 
			Config.GetString("strFN_CDBN_Settings");

		EnvVar.Inference_IterationAmount =
			Config.GetInt("Inference_IterationAmount");

		EnvVar.strFN_CDBN_Model =
			Config.GetString("strFN_CDBN_Model");
		EnvVar.Training_LayerIndex_Begin =
			Config.GetInt("Training_LayerIndex_Begin");
		EnvVar.Training_LayerIndex_End =
			Config.GetInt("Training_LayerIndex_End");
		EnvVar.Training_flag_ContinueCurrent =
			!!Config.GetInt("Training_flag_ContinueCurrent");
		EnvVar.Training_flag_DrawImageResponse =
			!!Config.GetInt("Training_flag_DrawImageResponse");
		EnvVar.Training_IterationAmount =
			Config.GetInt("Training_IterationAmount");
		EnvVar.Training_CheckFrequency =
			Config.GetDouble("Training_CheckFrequency");

		EnvVar.DrawScale_Model = Config.GetDouble(
			"DrawScale_Model", 1.0);
		EnvVar.DrawScale_Image = Config.GetDouble(
			"DrawScale_Image", 0.0);
		EnvVar.DrawResponse_LayerAmount = Config.GetInt(
			"DrawResponse_LayerAmount", 0);

		ofstream outFile_GlobalLog(strFN_OutLog.c_str());
		CExternalTools::ConsoleTools::MakeCurrentLine("", true);
		CExternalTools::ConsoleTools::OutputTextToStream(outFile_GlobalLog, "", true);
		ofstream outFile_ShortLog(strFN_ShortLog.c_str());

		// ********** Main Code Start **********

		if (1)
		{
			RBM_FLOAT InitialWeightValueVariance = 0.1f;
			string strPN_Maps = "maps\\";
			string strPN_Masks = "masks\\";
			string strPN_Models = "models\\";

			datasets_mnist::CImageSet MDS_Train;
			MDS_Train.myInit(EnvVar.strFN_MNIST_Image_Train, EnvVar.strFN_MNIST_Label_Train);

			CConvolutionalDBN_for_Image CDBN;
			CDBN.Settings.LoadFromFile(EnvVar.strFN_CDBN_Settings);
			CDBN.InitBySettings();

			if (op_TrainCDBN & 1)
			{
				CExternalTools::ConsoleTools::MakeDirectory(strPN_Maps, true);
				CExternalTools::ConsoleTools::MakeDirectory(strPN_Models, true);

				int src_label;
				CImg<int> imgSrc;
				CImg<int> imgOutput;
				CSimpleTypeArray<RBM_FLOAT> vSrcRealNumber;

				int i_train_layer;
				int i_train_iteration;
				int i_image, i_pixel;

				if (!(CDBN.LoadWeightModelFromFile(EnvVar.strFN_CDBN_Model)))
				{
					CDBN.InitBySettings();
					EnvVar.Training_LayerIndex_Begin = 0;
					EnvVar.Training_flag_ContinueCurrent = false;
					//CDBN.InitWeightModel(InitialWeightValueVariance);
					//EnvVar.Training_flag_ContinueCurrent = true;
				}
				if (EnvVar.Training_LayerIndex_End > CDBN.Depth)
					EnvVar.Training_LayerIndex_End = CDBN.Depth;


				int t_AmountImage;
				t_AmountImage = MDS_Train.AmountImage;
				//t_AmountImage = 1000;

				for ( i_train_layer = 0; i_train_layer < EnvVar.Training_LayerIndex_Begin; ++i_train_layer )
					CDBN.ComputeMaskImages(i_train_layer);
				for ( i_train_layer = EnvVar.Training_LayerIndex_Begin;
					i_train_layer < EnvVar.Training_LayerIndex_End; ++i_train_layer )
				{
					if (!EnvVar.Training_flag_ContinueCurrent)
						CDBN.GetLayer(i_train_layer).InitWeightItems(InitialWeightValueVariance);

					sprintf_s(strTemp, "%02d.input\\", i_train_layer);
					string strPN_Draw = strTemp;
					CExternalTools::ConsoleTools::MakeDirectory(strPN_Draw, true);

					CTimer TimerTotal;
					int AmountProcessedExample;

					CDBN.GetLayer(i_train_layer).ClearTimers();
					TimerTotal.Reset();
					AmountProcessedExample = 0;

					int NextCheck_i = 1;
					double NextCheck_f = 1.0;

					for ( i_train_iteration = 0; i_train_iteration < EnvVar.Training_IterationAmount; ++ i_train_iteration )
					{
						//CDBN.GetLayer(i_train_layer).ClearTimers();
						//TimerTotal.Reset(); TimerTotal.Start();
						//AmountProcessedExample = 0;

						// Update
						if (1) // (i_train_iteration < EnvVar.Training_UpdateStepAmount)
						{
							CDBN.ClearTrainingInfo(i_train_layer);
							TimerTotal.Start();

							for ( i_image = 0; i_image < t_AmountImage; ++i_image)
							{
								++AmountProcessedExample;

								src_label = MDS_Train.GetImage(i_image, imgSrc);
								vSrcRealNumber.Resize(MDS_Train.WH);
								for ( i_pixel = 0; i_pixel < MDS_Train.WH; ++i_pixel )
									vSrcRealNumber.pElement[i_pixel] = imgSrc.data[i_pixel] / 255.0f;

								CDBN.InitByInputSize(MDS_Train.W, MDS_Train.H);
								CDBN.SetVisibleState(0, vSrcRealNumber.pElement);

								CDBN.InferenceSubNetwork(
									0, i_train_layer, true, true, EnvVar.Inference_IterationAmount);
								if (i_train_iteration==0 && i_image==0)
									CDBN.PrepareTrainingLayer(i_train_layer);
								//CDBN.InferenceTimes_SubNetwork_BottomUp(
								//	0, i_train_layer, EnvVar.Inference_IterationAmount, true);

								// View Top Unit States
								if (1 && (i_train_iteration==0) && EnvVar.Training_flag_DrawImageResponse)
								{
									CDBN.DrawAllUnitsPropertyImage(
										i_train_layer, CConvolutionalDBN_for_Image::DRAW_UNITS_PROPERTY_MODE_VisibleState, (RBM_FLOAT)EnvVar.DrawScale_Image, 8, imgOutput);
									sprintf_s(strTemp, "%s%06d.jpg", strPN_Draw.c_str(), i_image);
									CExternalTools::mySaveImage(strTemp, imgOutput);
								} // View Top Unit States

								CDBN.ComputeBothGradient(i_train_layer, 1.0);
								CDBN.UpdateBothGradient(i_train_layer);
								if (i_image==0 || i_image==t_AmountImage-1 || CMyRandom::NextInt(t_AmountImage/1000)==0)
								{
									sprintf_s(strTemp, "Training [%d/%d, %d/%d, Update Step, %d/%d]...",
										i_train_layer, CDBN.Depth,
										i_train_iteration, EnvVar.Training_IterationAmount,
										i_image, t_AmountImage);
									CExternalTools::ConsoleTools::MakeCurrentLine(strTemp, false);
								}
							}

							TimerTotal.Stop();
						} // Update Step

						// Check & Print
						if (i_train_iteration==0
							|| i_train_iteration==EnvVar.Training_IterationAmount-1
							|| i_train_iteration+1==NextCheck_i
							//|| CMyRandom::NextInt(EnvVar.Training_CheckFrequency)==0
							)
						{
							if (EnvVar.Training_CheckFrequency < 1.01)
								++NextCheck_i;
							else
							{
								for ( ; ; )
								{
									NextCheck_f *= EnvVar.Training_CheckFrequency;
									NextCheck_i = CMyMath::TranslateNumber_Round(NextCheck_f, (int)0);
									if (NextCheck_i > i_train_iteration+1) break;
								}
							}

							TimerTotal.Start();

							// Check
							CDBN.ClearTrainingInfo(i_train_layer);
							for ( i_image = 0; i_image < t_AmountImage; ++i_image)
							{
								++AmountProcessedExample;

								src_label = MDS_Train.GetImage(i_image, imgSrc);
								vSrcRealNumber.Resize(MDS_Train.WH);
								for ( i_pixel = 0; i_pixel < MDS_Train.WH; ++i_pixel )
									vSrcRealNumber.pElement[i_pixel] = imgSrc.data[i_pixel] / 255.0f;

								CDBN.InitByInputSize(MDS_Train.W, MDS_Train.H);
								CDBN.SetVisibleState(0, vSrcRealNumber.pElement);

								CDBN.InferenceSubNetwork(
									0, i_train_layer, true, true, EnvVar.Inference_IterationAmount);
								//CDBN.InferenceTimes_SubNetwork_BottomUp(
								//	0, i_train_layer, EnvVar.Inference_IterationAmount, true);

								CDBN.ComputeBothGradient(i_train_layer, 1.0);
								if (i_image==0 || i_image==t_AmountImage-1 || CMyRandom::NextInt(t_AmountImage/1000)==0)
								{
									sprintf_s(strTemp, "Training [%d/%d, %d/%d, Check Step, %d/%d]...",
										i_train_layer, CDBN.Depth,
										i_train_iteration, EnvVar.Training_IterationAmount,
										i_image, t_AmountImage);
									CExternalTools::ConsoleTools::MakeCurrentLine(strTemp, false);
								}
							}
							CDBN.UpdateBothGradient(i_train_layer);

							TimerTotal.Stop();

							// Print
							RBM_FLOAT ve, hp, tot;
							tot = CDBN.GetCumulativeCosts(i_train_layer, ve, hp);
							CRBM_with_MaxPooling& tLM = CDBN.GetLayer(i_train_layer);
							sprintf_s(strTemp, "%d\t%d\t%.6lf\t%.6lf\t%.6lf\t%.6lf\t%.4lf\t%.4lf\t%.4lf",
								i_train_layer, i_train_iteration+1,
								(double)ve, (double)hp,
								(double)tLM.Training_DeltaWeight_ConDiv,
								(double)tLM.Training_DeltaWeight_SpaCod,
								(double)tLM.TimerInference.GetCumulativeTime()/(AmountProcessedExample+EPS),
								(double)tLM.TimerOther_ConDov.GetCumulativeTime()/(AmountProcessedExample+EPS),
								(double)TimerTotal.GetCumulativeTime()/(AmountProcessedExample+EPS));
							CExternalTools::ConsoleTools::OutputTextToStream(outFile_GlobalLog, strTemp, true);
							if (i_train_iteration == EnvVar.Training_IterationAmount-1)
								CExternalTools::ConsoleTools::MakeCurrentLine(strTemp, true);

						} // Check Step

						// Draw Model
						int i_v, i_k_h;
						CImg<int> imgY;
						//sprintf_s(strTemp, "%s%02d.%03d\\",
						//	strPN_Maps.c_str(), i_train_layer, i_train_iteration);
						//string strPN_CurrentMaps = strTemp;
						//CExternalTools::ConsoleTools::MakeDirectory(strPN_CurrentMaps, true);

						CDBN.ComputeMaskImages(i_train_layer);
						CConvolutionalDBN_for_Image::CLayerMaskImages& tMI = CDBN.GetSingleLayerMaskImages_Hidden(i_train_layer);

						int xDraw, yDraw;
						int wDraw, hDraw;
						for ( i_k_h = 0; i_k_h < tMI.U; ++i_k_h )
						{
							for ( i_v = 0; i_v < tMI.V; ++i_v )
							{
								CDBN.DrawMaskImage(
									i_train_layer, i_v, i_k_h, (RBM_FLOAT)EnvVar.DrawScale_Model, imgY);
								if (i_k_h==0 && i_v==0)
								{
									wDraw = imgY.width * tMI.V + 2;
									hDraw = imgY.height + 2;
									xDraw = CMyMath::TranslateNumber_Ceil(sqrt(1.0 * tMI.U / wDraw * hDraw));
									yDraw = (tMI.U + xDraw - 1) / xDraw;
									imgOutput.assign(xDraw * wDraw, yDraw * hDraw, 1, 1);
									imgOutput.fill(0);
								}
								imgOutput.draw_image(imgY, 
									i_k_h%xDraw*wDraw + imgY.width*i_v,
									i_k_h/xDraw*hDraw);
								//sprintf_s(strTemp, "%s%02d.%02d.jpg",
								//	strPN_CurrentMaps.c_str(), i_k_h, i_v);
								//CExternalTools::mySaveImage(strTemp, imgY);
							}
							sprintf_s(strTemp, "%s%02d.%03d.jpg",
								strPN_Maps.c_str(), i_train_layer, i_train_iteration);
							CExternalTools::mySaveImage(strTemp, imgOutput);
						}

						// Save Model
						sprintf_s(strTemp, "%s%02d.%03d.txt",
							strPN_Models.c_str(), i_train_layer, i_train_iteration);
						CDBN.SaveWeightModelToFile(strTemp);

					} // i_train_iteration
					CDBN.SaveWeightModelToFile(EnvVar.strFN_CDBN_Model);
				} // i_train_layer
			} // (op_TrainCDBN & 1)

			if (op_DrawCDBN & 1)
			{
				CExternalTools::ConsoleTools::MakeDirectory(strPN_Masks, true);

				CDBN.LoadWeightModelFromFile(EnvVar.strFN_CDBN_Model);
				int i_draw_layer, i_v, i_k_h;
				CImg<int> imgY;

				for ( i_draw_layer = 0; i_draw_layer < CDBN.Depth; ++i_draw_layer )
				{
					CDBN.ComputeMaskImages(i_draw_layer);
					CConvolutionalDBN_for_Image::CLayerMaskImages& tMI = CDBN.GetSingleLayerMaskImages_Hidden(i_draw_layer);

					for ( i_k_h = 0; i_k_h < tMI.U; ++i_k_h )
					{
						for ( i_v = 0; i_v < tMI.V; ++i_v )
						{
							CDBN.DrawMaskImage(
								i_draw_layer, i_v, i_k_h, (RBM_FLOAT)EnvVar.DrawScale_Model, imgY);
							sprintf_s(strTemp, "%s%02d.%02d.%02d.jpg",
								strPN_Masks.c_str(), i_draw_layer, i_k_h, i_v);
							CExternalTools::mySaveImage(strTemp, imgY);
						}
					}
				}
			} // (op_DrawCDBN & 1)

			if (op_DrawCDBN & 2)
			{
				int src_label;
				CImg<int> imgSrc;
				CImg<int> imgOutput;
				CSimpleTypeArray<RBM_FLOAT> vSrcRealNumber;

				int i_draw_layer;
				int i_image, i_pixel;

				if (!(CDBN.LoadWeightModelFromFile(EnvVar.strFN_CDBN_Model)))
				{
					CDBN.InitBySettings();
					CDBN.InitWeightModel(InitialWeightValueVariance);
					EnvVar.DrawResponse_LayerAmount = 0;
				}
				if (EnvVar.DrawResponse_LayerAmount > CDBN.Depth)
					EnvVar.DrawResponse_LayerAmount = CDBN.Depth;

				CDBN.InitByInputSize(MDS_Train.W, MDS_Train.H);

				int t_AmountImage;
				t_AmountImage = MDS_Train.AmountImage;
				t_AmountImage = 100;

				CSimpleTypeArray<string> List_strPN_Draw;
				List_strPN_Draw.Clear();
				for ( i_draw_layer = 0; i_draw_layer < EnvVar.DrawResponse_LayerAmount; ++i_draw_layer )
				{
					sprintf_s(strTemp, "%02d.output\\", i_draw_layer);
					string strPN_Draw = strTemp;
					CExternalTools::ConsoleTools::MakeDirectory(strPN_Draw, true);
					List_strPN_Draw.PushBack(strPN_Draw);
				}

				for ( i_image = 0; i_image < t_AmountImage; ++i_image)
				{
					src_label = MDS_Train.GetImage(i_image, imgSrc);
					vSrcRealNumber.Resize(MDS_Train.WH);
					for ( i_pixel = 0; i_pixel < MDS_Train.WH; ++i_pixel )
						vSrcRealNumber.pElement[i_pixel] = imgSrc.data[i_pixel] / 255.0f;
					CDBN.SetVisibleState(0, vSrcRealNumber.pElement);
					CDBN.InferenceTimes_SubNetwork_BottomUp(
						0, EnvVar.DrawResponse_LayerAmount, EnvVar.Inference_IterationAmount, true);

					for ( i_draw_layer = 0; i_draw_layer < EnvVar.DrawResponse_LayerAmount; ++i_draw_layer )
					{
						string strPN_Draw = List_strPN_Draw.pElement[i_draw_layer];
						CDBN.DrawAllUnitsPropertyImage(
							i_draw_layer, CConvolutionalDBN_for_Image::DRAW_UNITS_PROPERTY_MODE_HiddenExpectation, (RBM_FLOAT)EnvVar.DrawScale_Image, 8, imgOutput);
						sprintf_s(strTemp, "%s%d.%06d.jpg", strPN_Draw.c_str(), src_label, i_image);
						CExternalTools::mySaveImage(strTemp, imgOutput);
					} // i_draw_layer

					if (CMyRandom::NextInt(t_AmountImage/1000)==0)
					{
						sprintf_s(strTemp, "Drawing output of first [%d] layers: image [%d/%d]...",
							EnvVar.DrawResponse_LayerAmount,
							i_image, t_AmountImage);
						CExternalTools::ConsoleTools::MakeCurrentLine(strTemp, false);
					}
				} // i_image

				sprintf_s(strTemp, "Finished drawing output of first [%d] layers;", EnvVar.DrawResponse_LayerAmount);
				CExternalTools::ConsoleTools::OutputTextToStream(outFile_GlobalLog, strTemp, true);
				CExternalTools::ConsoleTools::MakeCurrentLine(strTemp, true);
			} // (op_DrawCDBN & 2)

		} // Main Code

		outFile_ShortLog.close();
		outFile_ShortLog.clear();
		CExternalTools::ConsoleTools::MakeCurrentLine("", true);
		CExternalTools::ConsoleTools::OutputTextToStream(outFile_GlobalLog, "", true);
		outFile_GlobalLog.close();
		outFile_GlobalLog.clear();
	}

	CImageEx::Release();
	return 0;
}

